"""
单变量线性回归
fθ(x)=θ0+θ1*x1
"""
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from linear_regression import LinearRegression
import matplotlib

def get_train_test_data(data,feature_column_name,label_column_name,frac=0.8):
    # 随机抽取80%的数据 一共155条数据  抽取124条数据
    train_data: pd.DataFrame = data.sample(frac=0.8)
    # 31 条数据
    test_data: pd.DataFrame = data.drop(train_data.index)
    return train_data,test_data
if __name__ == '__main__':
    data: pd.DataFrame = pd.read_csv('../data/world-happiness-report-2017.csv')
    feature_column_name = ['Economy..GDP.per.Capita.']
    label_column_name = ['Happiness.Score']
    (train_data, test_data)=get_train_test_data(data,feature_column_name=feature_column_name,
                                              label_column_name=label_column_name)
    #
    x_train: np.ndarray = train_data[feature_column_name].values
    y_train: np.ndarray = train_data[label_column_name].values
    #
    x_test: np.ndarray = test_data[feature_column_name].values
    y_test: np.ndarray = test_data[label_column_name].values
    #
    plt.scatter(x_train, y_train, label='Train data')
    plt.scatter(x_test, y_test, label='Test data')
    plt.xlabel(feature_column_name)
    plt.ylabel(label_column_name)
    plt.title('Happy')
    plt.legend()
    plt.show()
    #
    num_iters=500
    learn_rate=0.01
    linear_regression:LinearRegression = LinearRegression(x_train, y_train)
    theta, cost_history = linear_regression.train(learn_rate, num_iterations=num_iters)
    print('开始时的损失：', cost_history[0])
    print('训练后的损失：', cost_history[-1])
    # 画图
    plt.plot(range(num_iters), cost_history)
    plt.xlabel('Iter')
    plt.ylabel('cost')
    plt.title('GD')
    plt.show()

    predictions_num = 100
    x_predictions = np.linspace(x_train.min(), x_train.max(), predictions_num).reshape(predictions_num, 1)
    y_predictions = linear_regression.predict(x_predictions)

    plt.scatter(x_train, y_train, label='Train data')
    plt.scatter(x_test, y_test, label='Test data')
    plt.plot(x_predictions, y_predictions, 'r', label='Prediction')
    plt.xlabel(feature_column_name)
    plt.ylabel(label_column_name)
    plt.title('Happy')
    plt.legend()
    plt.show()